{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.Teleoprate the Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tutorial 4 is the example usage of teleoperation by human-expert."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from PIL import Image\n",
    "from VLABench.envs import load_env\n",
    "from VLABench.robots import *\n",
    "from VLABench.tasks import *\n",
    "from VLABench.configs import name2config\n",
    "from VLABench.utils.utils import find_key_by_value\n",
    "from VLABench.utils.interface import Interface\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 initialize env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TimeStep(step_type=<StepType.FIRST: 0>, reward=None, discount=None, observation={})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "robot = \"franka\"\n",
    "task = \"take_chemistry_experiment\"\n",
    "with open(os.path.join(os.getenv(\"VLABENCH_ROOT\"), \"configs/task_config.json\"), \"r\") as f:\n",
    "    TASK_CONFIG = json.load(f)\n",
    "task_series = find_key_by_value(name2config, task)\n",
    "task_config = TASK_CONFIG.get(task_series)\n",
    "default_config = TASK_CONFIG.get(\"default\")\n",
    "default_config.update(task_config)\n",
    "env = load_env(task, robot=robot, config=default_config)\n",
    "env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Load the interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================== Operation Menu ==============================\n",
      "Press 'w', 's', 'a', 'd', 'i', 'k' to move the robot along y, x, z axis.\n",
      "Press ',', '.' to rotate the robot along raw axis.\n",
      "Press 'n', 'm' to rotate the robot along pitch axis.\n",
      "Press 'j', 'l' to rotate the robot along yaw axis.\n",
      "Press 'space' to grasp.\n",
      "Press 'backspace' to delete the last step.\n",
      "Press 'enter' to reset.\n",
      "Press 'esc' to quit.\n",
      "================================================================================\n",
      "============================== Operation Menu ==============================\n",
      "Press 'w', 's', 'a', 'd', 'i', 'k' to move the robot along y, x, z axis.\n",
      "Press ',', '.' to rotate the robot along raw axis.\n",
      "Press 'n', 'm' to rotate the robot along pitch axis.\n",
      "Press 'j', 'l' to rotate the robot along yaw axis.\n",
      "Press 'space' to grasp.\n",
      "Press 'backspace' to delete the last step.\n",
      "Press 'enter' to reset.\n",
      "Press 'esc' to quit.\n",
      "================================================================================\n",
      "Quit the interface\n",
      "time cost:  57.867743253707886\n"
     ]
    }
   ],
   "source": [
    "render_options = {\n",
    "    \"width\": 480,\n",
    "    \"height\": 480\n",
    "}\n",
    "save_dir = \"/media/shiduo/LENOVO_USB_HDD/dataset/interface_test\"\n",
    "\n",
    "interface = Interface(env, \n",
    "                      save_dir=save_dir,\n",
    "                      dir_name=\"debug\",\n",
    "                      render_options=render_options)\n",
    "filename = time.strftime(\"%Y-%m-%d-%H-%M-%S\", time.localtime())\n",
    "interface.reset()\n",
    "start_time = time.time()\n",
    "while True:\n",
    "    interface.render(camera_ids=[0,1,2,3], \n",
    "                        column=2, \n",
    "                        render_only=False)\n",
    "    if env.task.should_terminate_episode(env.physics):\n",
    "        interface.save_single_data(index=0, filename=filename)\n",
    "        break\n",
    "    if interface.quit:\n",
    "        print(\"Quit the interface\")\n",
    "        break\n",
    "interface.close()\n",
    "print(\"time cost: \", time.time()-start_time)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vlabench",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
